import json

from llama_index.core.agent.workflow import AgentWorkflow

from llamaIndex.CommonClient import llm
from llama_index.core.workflow import Context
from llama_index.core.workflow import JsonPickleSerializer, JsonSerializer
# 维护状态

async def set_name(ctx: Context, name: str) -> str:
    state = await ctx.store.get("state")
    state["name"] = name
    await ctx.store.set("state", state)
    return f"Name set to {name}"

workflow = AgentWorkflow.from_tools_or_functions(
    [set_name],
    llm=llm,
    system_prompt="You are a helpful assistant that can set a name.",
    initial_state={"name": "unset"},
)
ctx = Context(workflow)
async def main():


    # check if it knows a name before setting it
    response = await workflow.run(user_msg="What's my name?", ctx=ctx)
    print(str(response))

    # set the name using a tool
    response2 = await workflow.run(user_msg="My name is Laurie", ctx=ctx)
    print(str(response2))
    # retrieve the value from the state directly
    state = await ctx.store.get("state")
    print("Name as stored in state: ", state["name"])

    #序列化
    ctx_dict = ctx.to_dict(serializer=JsonSerializer())
    print(ctx_dict)
    # 保存到文件
    with open("context.json", "w") as f:
        json.dump(ctx_dict, f)
    #用反序列化的上下文区问
    #反序列化报错，此方法不可用，具体还解决不了
    # restored_ctx = Context.from_dict(
    #     workflow, ctx_dict, serializer=JsonSerializer()
    # )


    response3 = await workflow.run(user_msg="我叫什么名字?", ctx=ctx)
    print(str(response3))

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())